# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
import random
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


@GENERATOR_REGISTRY.register("ascend_generate_moe_token_unpermute_with_routing_map")
class MoeTokenUnpermuteWithRoutingMapGenerator(CaseGenerator):

    def __init__(self, config):
        super().__init__(config)
        self.tensor_dim = 0
        self.range_is_null = False

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        tokens_num = random.randint(1,4096)
        hidden_size = random.randint(1,7168)
        experts_num = random.randint(16,256)
        top_k = random.randint(1,16)
        capacity = random.randint(1,16)

        #permuted_tokens [experts_num*capacity, hidden_size] #0
        case_config.inputs[0].shape = [experts_num*capacity, hidden_size]
        #sorted_indices [experts_num * capacity] #1 
        case_config.inputs[1].shape = [experts_num * capacity]
        #routing_map [tokens_num, experts_num] #2
        case_config.inputs[2].shape = [tokens_num, experts_num]
        #probs [tokens_num, experts_num] #3
        case_config.inputs[3].shape = [tokens_num, experts_num]

        #permuted_tokens 和 probs的数据类型一致
        x_dtype = case_config.inputs[0].dtype
        case_config.inputs[3].dtype = x_dtype

        # #attr 赋值
        # restoreShapeOptional (tokens_num, hidden_size) #5
        case_config.inputs[5][0].range_values = tokens_num
        case_config.inputs[5][1].range_values = hidden_size

        return case_config